// com/tidu/strategy/task/ParallelWriteTask.java
package com.tidu.strategy.task;

import com.tidu.utils.DbUtilsTemplate;
import lombok.extern.slf4j.Slf4j;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Callable;
import java.util.stream.Collectors;

@Slf4j
public class ParallelWriteTask implements Callable<Integer> {
    private final List<Map<String, Object>> dataBatch;
    private final DbUtilsTemplate targetDbUtils;
    private final String tableName;
    private final List<String> columns;

    public ParallelWriteTask(List<Map<String, Object>> dataBatch, DbUtilsTemplate targetDbUtils, String tableName, List<String> columns) {
        this.dataBatch = dataBatch;
        this.targetDbUtils = targetDbUtils;
        this.tableName = tableName;
        this.columns = columns;
    }

    @Override
    public Integer call() throws Exception {
        if (dataBatch.isEmpty()) {
            return 0;
        }

        Connection conn = null;
        try {
            conn = targetDbUtils.getConnection();
            conn.setAutoCommit(false);

            // 构建UPSERT SQL
            String insertSql = buildUpsertSQL();

            try (PreparedStatement ps = conn.prepareStatement(insertSql)) {
                // 优化的批处理
                for (int i = 0; i < dataBatch.size(); i++) {
                    Map<String, Object> map = dataBatch.get(i);
                    for (int j = 0; j < columns.size(); j++) {
                        Object value = map.get(columns.get(j));
                        ps.setObject(j + 1, value);
                    }
                    ps.addBatch();

                    // 每5000条记录执行一次批处理
                    if ((i + 1) % 5000 == 0) {
                        ps.executeBatch();
                        ps.clearBatch();
                    }
                }

                // 执行剩余的批处理
                ps.executeBatch();
                ps.clearBatch();
                conn.commit();

                log.debug("Parallel write task completed {} records", dataBatch.size());
                return dataBatch.size();
            }
        } catch (Exception e) {
            if (conn != null) {
                try {
                    conn.rollback();
                } catch (Exception rollbackEx) {
                    log.error("Failed to rollback in parallel task", rollbackEx);
                }
            }
            log.error("Error in parallel write task: {}", e.getMessage(), e);
            throw e;
        } finally {
            if (conn != null) {
                try {
                    conn.close();
                } catch (Exception e) {
                    log.warn("Failed to close connection in parallel task", e);
                }
            }
        }
    }

    private String buildUpsertSQL() {
        StringBuilder sqlBuilder = new StringBuilder();
        sqlBuilder.append("INSERT INTO ").append(tableName).append(" (").append(String.join(",", columns)).append(") VALUES (")
                .append(columns.stream().map(k -> "?").collect(Collectors.joining(","))).append(") ON DUPLICATE KEY UPDATE ");

        // 构建更新部分，排除主键字段
        String updateClause = columns.stream()
                .filter(column -> !column.equals("id"))
                .map(column -> column + " = VALUES(" + column + ")")
                .collect(Collectors.joining(", "));

        // 如果没有非主键字段可更新，则至少更新一个字段为自身值
        if (updateClause.isEmpty()) {
            updateClause = "id = VALUES(id)";
        }

        sqlBuilder.append(updateClause);
        return sqlBuilder.toString();
    }
}
